#coding=utf8
import re, socket, hashlib, base64, struct
from struct import pack
from utils.jsonUtils import  obj_to_str
from ws_setting import core_logger

puri = re.compile("/?([^/]+)(.*)")

#versions = ["WebSocketProtocolV00","WebSocketProtocolV13"]
#handshake_providers = [import_module(versions)]


class WebSocketProtocol():
    def forbidConnection(self, message):
        message = "HTTP/1.1 403 Forbidden\r\nContent-Length: " + str(len(message)) + "\r\n\r\n" + message
        self.socket.send(message.encode())

    #        return self.transport.loseConnection(message)
    def handshake(self):
        pass

    def validate(self):
        pass

    def send_data(self, message):
        pass

    def recv_data(self, count=256):
        pass

    def __eq__(self, obj):
        if hasattr(obj, "socket") and hasattr(self, "socket"):
            return self.socket == obj.socket
        return False

patternSec1 = re.compile("Sec-WebSocket-Key1: (.*)\r\n")
patternSec2 = re.compile("Sec-WebSocket-Key2: (.*)\r\n")

class WebSocketProtocolV00(WebSocketProtocol):
    def __init__(self, header_info, socket):
        self.header_info = header_info
        self.socket = socket

    def handshake(self):
        core_logger.info("using ws prtocol 00")
        _header = self.header_info

        c = lambda index: _header.token[index]

        challenge = pack('>II8B',
            self.part(_header.sec_websocket_key1), self.part(_header.sec_websocket_key2),
            c(0), c(1), c(2), c(3), c(4), c(5), c(6), c(7))
        _hash = hashlib.md5(challenge).digest()
        our_handshake = []
        our_handshake.append("HTTP/1.1 101 Web Socket Protocol Handshake")
        our_handshake.append("Upgrade: WebSocket")
        our_handshake.append("Connection: Upgrade")
        our_handshake.append("Sec-WebSocket-Origin: %s" % _header.origin)
        our_handshake.append("Sec-WebSocket-Location: ws://%s%s\r\n\r\n" % (_header.host, _header.uri))

        response = "\r\n".join(our_handshake)

        return (response.encode('utf-8') + _hash)


    def validate(self):
        vali = lambda key: hasattr(self.header_info, key)
        return vali("sec_websocket_key1") and vali("sec_websocket_key2")

    def getHeaders(self):
        _bytes = self.data[len(self.data) - 8:]

        key1 = patternSec1.findall(self.data)
        key2 = patternSec2.findall(self.data)

        return [key1[0], key2[0], _bytes]

    def part(self, token):
        digits = "".join(re.compile('\d').findall(token))
        count = token.count(" ")
        return int(int(digits) / count)

    def send_data(self, data):
        """
            send data to socket 
            data ：the data will be send
            clients ：接受数据的客户端, 如果clients为None，则将数据发送给当前用户
        """
        response_data = obj_to_str(data)
        if not type(response_data) == bytearray:
            if type(response_data) == str:
                response_data = response_data.encode('utf-8') 
                if type(response_data) == str:
                    response_data = b"\x00"+bytearray(response_data)+b"\xff"
                else:
                    response_data = b"\x00" + response_data +  b"\xff"
            else:
                response_data = obj_to_str(response_data)

        if response_data:
            self.socket.send(response_data)
            core_logger.debug("send response_data %s" % response_data)
        else:
            core_logger.error("no data to send")

    def recv_data(self, count=128):
        finished = lambda data: not data or data[-1:] == b'\xff' or data[-1:] == b'\x00'
        #recevice socket data after handshake
        data = b''
        #get all data from socket
        core_logger.debug("received data")
        while True:
            b = self.socket.recv(count)
            data += b
            if finished(data):
                break

        if not data or data == b'\xff\x00':
            #self.close_socket()
            return None

        core_logger.debug("recv_data = %s " % data)

        return data[1:-1]


    def close_socket(self):
        try:
            #try to close connection, if this throw an exception means that the connection was closed by client
            self.socket.send(bytearray('\xff\x00'))
        except :
            pass
        self.socket.close()
        core_logger.debug("socket closed")

class WebSocketProtocolV13(WebSocketProtocol):
    def __init__(self, header_info, socket):
        self.header_info = header_info
        self.socket = socket
        self._message_buffer = b''
        self.rev_data = ""

    def handshake(self):
        core_logger.info("using ws prtocol13")

        _hash = self.validate_websocket_key(self.header_info.sec_websocket_key)

        our_handshake = []
        our_handshake.append("HTTP/1.1 101 Switching Protocols")
        our_handshake.append("Upgrade: WebSocket")
        our_handshake.append("Connection: Upgrade")
        #        our_handshake.append("WebSocket-Origin: " + self.header_info.origin)
        our_handshake.append("Sec-WebSocket-Accept: " + _hash)
        #        our_handshake.append("Sec-WebSocket-Location: ws://%s%s\r\n\r\n" % (headerInfo.host, headerInfo.uri))

        response = "\r\n".join(our_handshake) + "\r\n\r\n"

        return response.encode()


    def validate(self):
        return hasattr(self.header_info, "sec_websocket_version") and self.header_info.sec_websocket_version == "13"

    def validate_websocket_key(self, ws_key):
        gen_uuid = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
        _value = (ws_key + gen_uuid).encode()
        response_key = base64.b64encode(hashlib.sha1(_value).digest())
        return response_key.decode()

    def send_data(self, data):
        data = obj_to_str(data)
        length = len(data)
        newFrame = []
        newFrame.append(0x81)
        newFrame = bytearray(newFrame)
        if length <= 125:
            newFrame.append(length)
        elif length > 125 and length < 65536:
            newFrame.append(126)
            newFrame += struct.pack('!H', length)
        elif length >= 65536:
            newFrame.append(127)
            newFrame += struct.pack('!Q', length)

        newFrame += data.encode()
        self.socket.send(newFrame)

    def _init_frame_header(self):
        self._frame_fin = 0
        self._frame_rsv = 0
        self._frame_ops = 0
        self._frame_mask = False
        self._header_index = 0
        self._frame_payload_len = 0
        self._message_buffer = b''

    def recv_data(self, count = 16):
        #get all data from socket
        core_logger.debug("received data")

        self._init_frame_header()
        data = self.socket.recv(count)
        payload_data = self.rawDataReceived(bytearray(data))

        core_logger.debug("recev data is {}".format(payload_data))
        return payload_data

    def rawDataReceived(self, data):
        self._raw_data_len = len(data)

        self._processFrameHeader(data)

        remain_payload_len = (self._frame_header_len+ self._frame_payload_len) - self._raw_data_len
        if remain_payload_len > 0:
            data += self.socket.recv(remain_payload_len)
        self._message_buffer += self._extractMessageFromFrame(data)
        core_logger.debug("\n[ recv_data = \n {} \n]".format(self._message_buffer))
#        if (self._raw_data_len - self._header_index) < self._frame_payload_len:
#            self._partial_data = data
#            return


        if self._frame_payload_len == len(self._message_buffer):
            return self._message_buffer



    def _extractMessageFromFrame(self, data):
        i = self._header_index

        # when payload is masked, extract frame mask
        frame_mask = None
        frame_mask_array = []
        if self._frame_mask:
            frame_mask = data[i:i + 4]
            for j in range(0, 4):
                frame_mask_array.append(frame_mask[j])
            i += 4
            payload = bytearray(data[i:i + self._frame_payload_len])
            for k in range(0, self._frame_payload_len):
                payload[k] ^= frame_mask_array[k % 4]

            return payload

    def _processFrameHeader(self, data):
        # first byte contains fin, rsv and ops
        b = data[0]
        #fin: 
        #Indicates that this is the final fragment in a message.  The first
        #fragment MAY also be the final fragment
        self._frame_fin = (b & 0x80) != 0
        #rsv
        #MUST be 0 unless an extension is negotiated which defines meanings
        #for non-zero values.  If a nonzero value is received and none of
        #the negotiated extensions defines the meaning of such a nonzero
        #value, the receiving endpoint MUST _Fail the WebSocket Connection_.
        self._frame_rsv = (b & 0x70) >> 4
        self._frame_ops = b & 0x0f

        # second byte contains mask and payload length
        b = data[1]
        self._frame_mask = (b & 0x80) != 0

        frame_payload_len1 = b & 0x7f

        if (self._frame_mask):
            mask_len = 4
        else:
            mask_len = 0

        # i is frame index. It's at 2 here because we've already processed the
        # first 2 bytes of the frame.
        i = 2

        if frame_payload_len1 < 126:
            self._frame_header_len = i + mask_len
            self._frame_payload_len = frame_payload_len1
        elif frame_payload_len1 == 126:
            self._frame_header_len = i + 2 + mask_len
            self._frame_payload_len = struct.unpack("!H", data[i:i + 2])[0]
            i += 2
        elif frame_payload_len1 == 127:
            self._frame_header_len = i + 8 + mask_len
            self._frame_payload_len = struct.unpack("!Q", data[i:i + 8])[0]
            i += 8

        self._header_index = i

        core_logger.debug("_frame_fin = {}".format(self._frame_fin))
        core_logger.debug("_frame_rsv = {}".format(self._frame_rsv))
        core_logger.debug("_frame_ops = {}".format(self._frame_ops))
        core_logger.debug("_frame_mask = {}".format(self._frame_mask))
        core_logger.debug("_header_index = {}".format(self._header_index))
        core_logger.debug("_frame_payload_len = {}".format(self._frame_payload_len))

    def close_socket(self):
        self.socket.close()
        core_logger.debug("socket closed")


class HeaderInfo:
    def __init__(self, orgin_data=None):
        self.data = orgin_data;
        if orgin_data:
            cdata = orgin_data.decode('utf-8', 'ignore')
            headers = cdata.split("\r\n")
            self.scheme, self.uri, self.protocol = headers[0].split()

            self.init_connection_info(headers[1:])
            self.parameters = dict()
            if "?" in self.uri:
                self.full_url, paraStr = self.uri.split("?")
                a_list_of_queryStr = [pair.split('=', 1) for pair in paraStr.split('&')]
                self.parameters = dict(a_list_of_queryStr)

    def init_connection_info(self, additional_info):
        for info in additional_info:
            if info and ":" in info:
                key, val = info.split(":", 1)
                info_key = key.lower().strip().replace("-", "_")
                self.__dict__[info_key] = val.strip()
            else:
                self.token = bytearray(self.data[-8:])

    def parseUrl(self):
        if "?" in self.uri:
            return puri.findall(self.uri.split("?")[0])[0]
        else:
            #no parameter
            return (self.uri, "")


def handshake(socket):
    header_info = HeaderInfo(socket.recv(1024));
    for pmodule in handshake_providers:
        _provider = pmodule(header_info, socket)
        if _provider.validate():
            response_data = _provider.handshake()
            #send handshake response back to client 
            _provider.socket.send(response_data)
            break

    return _provider

handshake_providers = [WebSocketProtocolV00, WebSocketProtocolV13]
